package com.qin.algorithm;

public class RBTree<K, V> {
    private Node<K, V> root;
    private int size;
    
    /**
     * 查找
     */
    public V get(K k) {
        Node<K, V> node = getNode(k);
        return node == null ? null : node.value;
    }

    /**
     * 查找 Node节点
     */
    private Node<K, V> getNode(K key) {
        Node<K, V> p = root;
        @SuppressWarnings("unchecked")
        Comparable<? super K> k = (Comparable<? super K>) key;
        
        while(p != null) {
            int cmp = k.compareTo(p.key);
            if(cmp < 0) {
                p = p.left;
            } else if(cmp > 0) {
                p = p.right;
            } else {
                return p;
            }
        }
        
        return null;
    }

    /**
     * 存放键值对
     * @return true表示存放新增，false表示修改
     */
    public boolean put(K key, V value) {
        if(root == null) {
            root = new Node<>(key, value, null);
            root.color = Node.BLACK;
            size = 1;
            return true;
        }
        
        Node<K, V> t = root;
        Node<K, V> parent;
        
        int cmp;
        
        @SuppressWarnings("unchecked")
        Comparable<? super K> k = (Comparable<? super K>) key;
        do {
            parent = t;
            cmp = k.compareTo(t.key);
            if(cmp < 0) {
                t = t.left;
            } else if(cmp > 0) {
                t = t.right;
            } else {
                t.value = value;
                return false;
            }
            
        }while(t != null);

        Node<K, V> n = new Node<>(key, value, parent);
        
        if(cmp < 0) {
            parent.left = n;
        } else {
            parent.right = n;
        }

        fixAfterInsertion(n);
        
        ++size;
        
        return true;
    }

    /**
     * 维护插入后的规则
     */
    private void fixAfterInsertion(Node<K, V> x) {
        x.color = Node.RED;

        for(Node<K, V> xp, xpp, xppl, xppr; ;) {
            if((xp = x.parent) == null) { // case 1: x是根节点  
                x.color = Node.BLACK;
                return;
            } else if(xp.isBlack()) { // case 2: 父节点是黑色  
                return; // hong   
            } else if((xpp = xp.parent) == null) { // case 3: 父节点是root  
                xp.color = Node.BLACK;
                return;
            }
            if(xp == (xppl = xpp.left)) {
                if((xppr = xpp.right) != null && xppr.isRed()) {  // case 4: x、xp、叔父都是红色  
                    xp.color = Node.BLACK;
                    xppr.color = Node.BLACK;
                    xpp.color = Node.RED;
                    x = xpp; // 向上处理 xpp
                } else {
                    // 如果是RL，那么左转xp变成LL  
                    if (x == xp.right) {
                        rotateLeft(x = xp);
                        xpp = (xp = x.parent) == null ? null : xp.parent;
                    }
                    // case 5: 如果是LL，那么：xp变成黑色、xpp变红、右旋转xpp  
                    if (xp != null) {
                        xp.color = Node.BLACK;
                        if (xpp != null) {
                            xpp.color = Node.RED;
                            rotateRight(xpp);
                        }
                    }
                }
            } else {
                if(xppl != null && xppl.isRed()) {   // case 4: x、xp、叔父都是红色  
                    xp.color = Node.BLACK;
                    xppl.color = Node.BLACK;
                    xpp.color = Node.RED;
                    x = xpp; // 向上处理 xpp           
                } else {
                    // 如果是LR，那么右转xp变成RR  
                    if (x == xp.left) {
                        rotateRight(x = xp);
                        xpp = (xp = x.parent) == null ? null : xp.parent;
                    }
                    // case 6: 如果是RR，那么：xp变成黑色、xpp变红、左旋转xpp  
                    if (xp != null) {
                        xp.color = Node.BLACK;
                        if (xpp != null) {
                            xpp.color = Node.RED;
                            rotateLeft(xpp);
                        }
                    }
                }
            }
        }
    }


    /**
     * 删除一个元素
     * @return true表示删除成功，false表示未找到该节点
     */    
    public boolean remove(K key) {
        @SuppressWarnings("unchecked")
        Comparable<? super K> k = (Comparable<? super K>) key;
        int cmp;
        Node<K, V> t = root;

        while (t != null) {
            cmp = k.compareTo(t.key);
            if(cmp < 0) {
                t = t.left;
            } else if(cmp > 0) {
                t = t.right;
            } else {
                // 找到了就是这个节点
                break;
            }
        }

        return doRemove(t);
    }

    /**
     * 真正完成节点删除的方法
     */
    private boolean doRemove(Node<K, V> x) {
        if(x == null) {
            return false;
        }
        
        Node<K, V> xl, xr;
        // 如果有左、右节点
        if(x.left != null && (xr = x.right) != null) {
            // 找后继节点替换
            while(xr.left != null) {
                xr = xr.left;
            }
            
            // 此时xr便是后继节点交换值
            x.key = xr.key;
            x.value = xr.value;
            x = xr;
        }
        
        // 如果不是叶子节点，必然是黑-红连接。只需要替换，然后删除叶子即可
        if((xl = x.left) != null) {
            x.key = xl.key;
            x.value = xl.value;
            x = xl;
        } else if((xr = x.right) != null){
            x.key = xr.key;
            x.value = xr.value;
            x = xr;
        }
        
        // 此时的x必然是叶子节点，如果它不是红色，那么删除前尝试平衡它
        if(x.isBlack()) {
            fixAfterDeletion(x);
        }
        
        // 删除叶子节点x
        deleteEntry(x);
        --size;
        
        return true;
    }

    /**
     * 将叶子节点x从父节点上删除
     */
    private void deleteEntry(Node<K, V> x) {
        Node<K, V> xp = x.parent;
        if(xp == null) {
            root = null;
            return;
        }
        
        if(xp.left == x) {
            xp.left = null;
        } else {
            xp.right = null;
        }
        
        x.parent = x.left = x.right = null;
    }

    /**
     * 修复在删除后叶子节点x处的失衡
     */
    private void fixAfterDeletion(Node<K, V> x) {
        Node<K, V> xp, xs, xsl, xsr;
        while(x != null && x.isBlack()) {
            if((xp = x.parent) == null) {
                break;
            } else if(xp.right == x) {
                if((xs = xp.left) != null && xs.isRed()) { // 兄弟红转成兄弟黑
                    xs.color = Node.BLACK;
                    xp.color = Node.RED;
                    rotateRight(xp);
                    xs = (xp = x.parent) == null ? null : xp.left;
                }
                
                assert xs != null;
                
                if((xsl = xs.left) != null && xsl.isRed()) { // 兄弟左子红
                    xsl.color = Node.BLACK;
                    xs.color = xp.color;
                    xp.color = Node.BLACK;
                    rotateRight(xp);
                    return;
                } else if((xsr = xs.right) != null && xsr.isRed()) { // 兄弟右子红
                    xsr.color = xp.color;
                    xp.color = Node.BLACK;
                    rotateLeft(xs);
                    rotateRight(xp);
                    return;
                } else { // 兄弟子全黑
                    xs.color = Node.RED;
                    if(xp.isRed()) {
                        xp.color = Node.BLACK;
                        return;
                    }
                    
                    x = xp;
                }
            } else {    // 上面的镜像
                if((xs = xp.right) != null && xs.isRed()) { // 兄弟红转成兄弟黑
                    xs.color = Node.BLACK;
                    xp.color = Node.RED;
                    rotateLeft(xp);
                    xs = (xp = x.parent) == null ? null : xp.right;
                }

                assert xs != null;

                if((xsr = xs.right) != null && xsr.isRed()) { // 兄弟右子红
                    xsr.color = Node.BLACK;
                    xs.color = xp.color;
                    xp.color = Node.BLACK;
                    rotateLeft(xp);
                    return;
                } else if((xsl = xs.left) != null && xsl.isRed()) { // 兄弟左子红
                    xsl.color = xp.color;
                    xp.color = Node.BLACK;
                    rotateRight(xs);
                    rotateLeft(xp);
                    return;
                } else { // 兄弟子全黑
                    xs.color = Node.RED;
                    if(xp.isRed()) {
                        xp.color = Node.BLACK;
                        return;
                    }

                    x = xp;
                }
            }
        }
    }

    /**
     * 获取长度
     */
    public int getSize() {
        return size;
    }

    /**
     * 将p左旋
     */
    private void rotateLeft(Node<K, V> p) {
        if(p != null) {
            Node<K, V> r = p.right;
            p.right = r.left;
            if(r.left != null) {
                r.left.parent = p;
            }
            
            r.parent = p.parent;
            if(p.parent == null) {
                root = r;
            } else if(p.parent.left == p){
                p.parent.left = r;
            } else {
                p.parent.right = r;
            }
            r.left = p;
            p.parent = r;
        }
    }

    /**
     * 将p右旋
     */
    private void rotateRight(Node<K, V> p) {
        if(p != null) {
            Node<K, V> l = p.left;
            p.left = l.right;
            if(l.right != null) {
                l.right.parent = p;
            }

            l.parent = p.parent;
            if(p.parent == null) {
                root = l;
            } else if(p.parent.left == p){
                p.parent.left = l;
            } else {
                p.parent.right = l;
            }
            l.right = p;
            p.parent = l;
        }
    }

    static final class Node<K, V> {

        public static final boolean RED = true;
        public static final boolean BLACK = false;

        K key;
        V value;

        Node<K, V> left;
        Node<K, V> right;
        Node<K, V> parent;

        boolean color;

        public Node(K key, V value,  Node<K, V> parent) {
            this.key = key;
            this.value = value;
            this.parent = parent;
        }

        public boolean isRed() {
            return color == RED;
        }

        public boolean isBlack() {
            return color == BLACK;
        }

        @Override
        public String toString() { return key + "=" + value; }
    }
}